#include "ImageHelper.h"
#include "swf/SWFTag.h"
#include <QImage>
#include <QImageWriter>
#include <QBuffer>
#include "QtZlib/zlib.h"

TagDefineBits* ImageHelper::createDefineBits( QImage* image, bool highQuality )
{
	if (highQuality)
	{
		return createDefineBitsJPEG3(image);
	}
	else
	{
		if (image->hasAlphaChannel())
			return createTagDefineBitsLossless2(image);
		else
			return createTagDefineBitsLossless(image);
	}
}

TagDefineBitsJPEG2* ImageHelper::createDefineBitsJPEG2( QImage* image )
{
	return nullptr;
}

TagDefineBitsJPEG3* ImageHelper::createDefineBitsJPEG3( QImage* image )
{
	uchar* bits = image->bits();
	size_t imageSize = image->width() * image->height();
	TagDefineBitsJPEG3* tag = TagDefineBitsJPEG3::create();
	tag->code = DefineBitsJPEG3;

	char* alphaData = (char*)malloc(imageSize);
	int alphaIndex = 0;
	for (size_t i = 0; i < imageSize; ++i)
	{
		char a = bits[i * 4 + 3];
		//char r = bits[i + 2];
		//char g = bits[i + 1];
		//char b = bits[i + 0];
		alphaData[alphaIndex++] = a;
	}
	char* dest = (char*)malloc(imageSize);
	z_uLongf destLen;
	compress((z_Bytef*)dest, &destLen, (z_Bytef*)alphaData, (z_uLongf)imageSize);

	tag->alphaData.copy(dest, destLen);
	free(alphaData);
	free(dest);

	//convert to jpeg
	QImage jpegImage = image->convertToFormat(QImage::Format_RGB888);
	QBuffer buffer;
	QImageWriter writer(&buffer, "jpeg");
	writer.setQuality(100);
	writer.write(jpegImage);
	QByteArray arr = buffer.data();
	tag->imageData.copy(arr.data(), arr.size());

	return tag;
}

TagDefineBitsLossless* ImageHelper::createTagDefineBitsLossless( QImage* image )
{
	TagDefineBitsLossless* tag = TagDefineBitsLossless::create();
	tag->code = DefineBitsLossless;
	tag->BitmapFormat = 5;
	tag->BitmapWidth = image->width();
	tag->BitmapHeight = image->height();

	int pixelCount = image->width() * image->height();
	uchar* pixelData = (uchar*)malloc(pixelCount * 4);
	uchar* sourcePixelData = image->bits();

	for (int i = 0; i < pixelCount; ++i)
	{
		pixelData[i * 4 + 0] = 0;
		pixelData[i * 4 + 1] = sourcePixelData[i * 3 + 0];
		pixelData[i * 4 + 2] = sourcePixelData[i * 3 + 1];
		pixelData[i * 4 + 3] = sourcePixelData[i * 3 + 2];
	}

	z_uLongf compressSize = pixelCount * 4;
	uchar* compressData = (uchar*)malloc(compressSize);
	int ok = compress(compressData, &compressSize, pixelData, pixelCount * 4);
	assert(ok == Z_OK);
	tag->ZlibBitmapData.copy((char*)compressData, compressSize);

	free(pixelData);
	free(compressData);

	return tag;
}

TagDefineBitsLossless2* ImageHelper::createTagDefineBitsLossless2( QImage* image )
{
	TagDefineBitsLossless2* tag = TagDefineBitsLossless2::create();
	tag->code = DefineBitsLossless2;
	tag->BitmapFormat = 5;
	tag->BitmapWidth = image->width();
	tag->BitmapHeight = image->height();

	int pixelCount = image->width() * image->height();
	uchar* pixelData = (uchar*)malloc(pixelCount * 4);
	uchar* sourcePixelData = image->bits();

	for (int i = 0; i < pixelCount; ++i)
	{
		pixelData[i * 4 + 0] = sourcePixelData[i * 4 + 3];
		pixelData[i * 4 + 1] = sourcePixelData[i * 4 + 0];
		pixelData[i * 4 + 2] = sourcePixelData[i * 4 + 1];
		pixelData[i * 4 + 3] = sourcePixelData[i * 4 + 2];
	}

	z_uLongf compressSize = pixelCount * 4;
	uchar* compressData = (uchar*)malloc(compressSize);
	int ok = compress(compressData, &compressSize, pixelData, pixelCount * 4);
	assert(ok == Z_OK);
	tag->ZlibBitmapData.copy((char*)compressData, compressSize);

	free(pixelData);
	free(compressData);

	return tag;
}

QImage* convertDefineBitsJPEG2ToImage(TagDefineBitsJPEG2* tag)
{
	QImage* img = new QImage();
	img->loadFromData((uchar*)tag->imageData.getData(), tag->imageData.getLength());

	return img;
}

QImage* convertDefineBitsJPEG3ToImage(TagDefineBitsJPEG3* tag)
{
	QImage temp;
	temp.loadFromData((uchar*)tag->imageData.getData(), tag->imageData.getLength());
	QImage* img = new QImage(temp.convertToFormat(QImage::Format_RGBA8888));
	
	z_uLongf pixelCount = img->width() * img->height();
	uchar* alphaData = (uchar*)malloc(pixelCount);
	int ok = uncompress(alphaData, &pixelCount, (z_Byte*)tag->alphaData.getData(), tag->alphaData.getLength());
	assert(ok == Z_OK);
	for (size_t i = 0; i < pixelCount; ++i)
	{
		char a = alphaData[i];
		img->bits()[i * 4 + 3] = a;
	}
	free(alphaData);

	return img;
}

QImage* convertDefineBitsLosslessToImage( TagDefineBitsLossless* tag )
{
	z_uLongf imageDataSize = tag->BitmapWidth * tag->BitmapHeight;
	if (tag->BitmapFormat == 3) //RGB
		imageDataSize *= 24;
	else if (tag->BitmapFormat == 4) //PIX15
		imageDataSize *= 16;
	else if (tag->BitmapFormat == 5) //PIX24
		imageDataSize *= 32;

	char* imageData = (char*)malloc(imageDataSize);
	int result = uncompress((z_Bytef*)imageData, &imageDataSize, (z_Bytef*)tag->ZlibBitmapData.getData(), tag->ZlibBitmapData.getLength());
	assert(result == Z_OK);

	SWFInputStream stream(imageData, imageDataSize);

	char* colorTable = nullptr;
	UI16 paddedWidth = tag->BitmapWidth;
	//read color table
	if (tag->BitmapFormat == 3)
	{
		paddedWidth = ceil(tag->BitmapWidth / 4.0f) * 4;
		UI16 tableSize = tag->BitmapColorTableSize + 1;
		colorTable = (char*)malloc(tableSize * 3);
		stream.readBytes(colorTable, tableSize * 3);
	}

	char* pixelData = (char*)malloc(paddedWidth * tag->BitmapHeight * 24);
	int pixelIndex = 0;
	int dataLen = 0;

	for (UI16 h = 0; h < tag->BitmapHeight; ++h)
	{
		for (UI16 w = 0; w < paddedWidth; ++w)
		{
			if (tag->BitmapFormat == 3)
			{
				int idx = stream.readUI8() & 0xff;
				assert(idx < tag->BitmapColorTableSize + 1);
				pixelData[pixelIndex + 0] = colorTable[idx * 3 + 0];//R
				pixelData[pixelIndex + 1] = colorTable[idx * 3 + 1];//G
				pixelData[pixelIndex + 2] = colorTable[idx * 3 + 2];//B
			}
			else if (tag->BitmapFormat == 4)
			{
				assert(false);
				stream.readUB(1);
				pixelData[pixelIndex + 0] = stream.readUB(5);
				pixelData[pixelIndex + 1] = stream.readUB(5);
				pixelData[pixelIndex + 2] = stream.readUB(5);
				dataLen += 2;
			}
			else if (tag->BitmapFormat == 5)
			{
				stream.skip(1);
				pixelData[pixelIndex + 0] = stream.readUI8();
				pixelData[pixelIndex + 1] = stream.readUI8();
				pixelData[pixelIndex + 2] = stream.readUI8();
				dataLen += 4;
			}
			pixelIndex += 3;
		}
		while ((dataLen % 4) != 0)
		{
			dataLen++;
			stream.skip(1);
		}
	}
	QImage* img = new QImage(paddedWidth, tag->BitmapHeight, QImage::Format_RGB888);
	memcpy(img->bits(), pixelData, pixelIndex);
	free(imageData);
	free(pixelData);
	if (colorTable)
		free(colorTable);

	return img;
}

QImage* convertDefineBitsLosslessToImage( TagDefineBitsLossless2* tag )
{
	z_uLongf imageDataSize = tag->BitmapWidth * tag->BitmapHeight;
	if (tag->BitmapFormat == 3) //RGB
		imageDataSize *= 24;
	else if (tag->BitmapFormat == 4) //PIX15
		imageDataSize *= 16;
	else if (tag->BitmapFormat == 5) //PIX24
		imageDataSize *= 32;

	char* imageData = (char*)malloc(imageDataSize);
	int result = uncompress((z_Bytef*)imageData, &imageDataSize, (z_Bytef*)tag->ZlibBitmapData.getData(), tag->ZlibBitmapData.getLength());
	assert(result == Z_OK);

	SWFInputStream stream(imageData, imageDataSize);

	char* colorTable = nullptr;
	UI16 paddedWidth = tag->BitmapWidth;
	//read color table
	if (tag->BitmapFormat == 3)
	{
		paddedWidth = ceil(tag->BitmapWidth / 4.0f) * 4;
		UI16 tableSize = tag->BitmapColorTableSize + 1;
		colorTable = (char*)malloc(tableSize * 4);
		stream.readBytes(colorTable, tableSize * 4);
	}

	char* pixelData = (char*)malloc(paddedWidth * tag->BitmapHeight * 32);

	int pixelIndex = 0;
	for (UI16 h = 0; h < tag->BitmapHeight; ++h)
	{
		for (UI16 w = 0; w < paddedWidth; ++w)
		{
			if (tag->BitmapFormat == 3)
			{
				int idx = stream.readUI8() & 0xff;
				assert(idx < tag->BitmapColorTableSize + 1);
				pixelData[pixelIndex + 0] = colorTable[idx * 4 + 0];//A
				pixelData[pixelIndex + 1] = colorTable[idx * 4 + 1];//R
				pixelData[pixelIndex + 2] = colorTable[idx * 4 + 2];//G
				pixelData[pixelIndex + 3] = colorTable[idx * 4 + 3];//B
			}
			else if (tag->BitmapFormat == 4 || tag->BitmapFormat == 5)
			{
				pixelData[pixelIndex + 3] = stream.readUI8();
				pixelData[pixelIndex + 0] = stream.readUI8();
				pixelData[pixelIndex + 1] = stream.readUI8();
				pixelData[pixelIndex + 2] = stream.readUI8();
			}
			pixelIndex += 4;
		}
	}
	QImage* img = new QImage(paddedWidth, tag->BitmapHeight, QImage::Format_RGBA8888);
	memcpy(img->bits(), pixelData, pixelIndex);
	free(imageData);
	free(pixelData);
	if (colorTable)
		free(colorTable);

	return img;
}

QImage* ImageHelper::convertDefineBitsToImage( TagDefineBits* tag )
{
	QImage* image = nullptr;
	switch (tag->code)
	{
	case DefineBitsLossless:
		image = convertDefineBitsLosslessToImage(dynamic_cast<TagDefineBitsLossless*>(tag));
		break;
	case DefineBitsLossless2:
		image = convertDefineBitsLosslessToImage(dynamic_cast<TagDefineBitsLossless2*>(tag));
		break;
	case DefineBitsJPEG2:
		image = convertDefineBitsJPEG2ToImage(dynamic_cast<TagDefineBitsJPEG2*>(tag));
		break;
	case DefineBitsJPEG3:
		image = convertDefineBitsJPEG3ToImage(dynamic_cast<TagDefineBitsJPEG3*>(tag));
		break;
	default:
		assert(false);
		break;
	}
	return image;
}
